#include "common.h"
#include "auxiliary.h"
#include "bregman.h"

Matcp get_fourier_mask(int rows, int cols, rltype sampleprob);


int main()
{
	// setup display windows
	namedWindow("Input image", 1);
	namedWindow("Fourier magnitude", 1);
	namedWindow("Reconstruction", 1);

	// split bregman constants
	const rltype mu = 10.0;
	const rltype lambda = 10.0;
	const rltype gamma = 0.1;
	const int nInner = 10;
	const int nOuter = 20;
	const rltype sampleprob = 0.15;

	// Read input data, and its identical real matrix for easy viewing
	Matcp imin = read_image("pr3ut.png");
	Matrl iminabs = cvabs(imin);
	rltype M = rltype(imin.rows);
	rltype N = rltype(imin.cols);

	
	// Create full 2D DFT of imin
	Matcp iminF;
	dft(imin, iminF);
	iminF = iminF/sqrt(M*N);

	// Get Fourier magnitude for sanity-check viewing
	Matrl iminFabs = cvabs(iminF);

	// Create Fourier sampling mask, and its identical real matrix for easy viewing
	Matcp R = get_fourier_mask(int(M), int(N), sampleprob);

	// Create measurement data as R.*iminF
	Matcp datain = elmult(R, iminF);

	// Construct u
	Matcp u = split_bregman(R, datain, mu, lambda, gamma, nInner, nOuter);

	// Take abs of u to get final output
	Matrl uabs = cvabs(u);

	// Display images
	imshow("Input image", iminabs);
	imshow("Fourier magnitude", iminFabs);
	imshow("Reconstruction", uabs);

	// Wait for a key input before closing the program
	waitKey(0);
	return 0;

}

// get_fourier_mask: create a (rows x cols) binary complex array, where each element is 1 with probability
// sampleprob and 0 otherwise
Matcp get_fourier_mask(int rows, int cols, rltype sampleprob)
{
	Matrl fmask = Mat::zeros(rows, cols, CV_RL);
	Matcp fmaskcp;
	RNG rng;
	rng.fill(fmask, RNG::UNIFORM, 0, 1);
	fmask = thresholdbi64(fmask, sampleprob, 1.0);
	fmask.at<rltype>(0) = 1.0; // The origin in Fourier space must be in the mask, for TV reconstruction.
	fmaskcp = merge2(fmask, Mat::zeros(rows, cols, CV_RL));
	return fmaskcp;
}